{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Map-Reduce"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plan\n",
    "\n",
    "1. **Ask the user** what kind of companies they want to invest in and how many options they want to see.\n",
    "2. **Find some matching stocks** based on the user’s interests.\n",
    "3. **In parallel**, gather key information about each stock — like financial reports, market performance, etc. (*This is the “map” step.*)\n",
    "4. **Analyze all the results** and create a final list that (*This is the “reduce” step.*):\n",
    "- Describes each stock\n",
    "- Explains why it’s a good investment\n",
    "- Ranks them by priority"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing_extensions import TypedDict\n",
    "from langchain_openai import ChatOpenAI\n",
    "import operator\n",
    "from typing import Annotated\n",
    "\n",
    "class InvestmentAdvisorState(TypedDict):\n",
    "    financial_area: str\n",
    "    stock_number: int\n",
    "    stock_tickers: list[str]\n",
    "    stock_details: Annotated[list, operator.add]\n",
    "    recommendation: str\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate list of stocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "\n",
    "class StockTickers(BaseModel):\n",
    "    stock_tickers: list[str]\n",
    "\n",
    "def generate_list_of_stocks(state: InvestmentAdvisorState):\n",
    "    prompt = f\"\"\"\n",
    "    You are an experienced financial analyst. \n",
    "    Based on current market trends and public data, suggest the top {state[\"stock_number\"]} publicly traded companies in the area of \"{state[\"financial_area\"]}\".\n",
    "\n",
    "    Return only a list of their stock tickers (symbols), without any explanations.\n",
    "    \"\"\"\n",
    "    \n",
    "    response = llm.with_structured_output(StockTickers).invoke(prompt)\n",
    "    return {\"stock_tickers\": response.stock_tickers}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_list_of_stocks({\n",
    "   \"financial_area\": \"small AI companies\",\n",
    "   \"stock_number\": 5\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fetch Stock Details (Map Step)\n",
    "\n",
    "see that we do not have to follow here overall graph state structure!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yfinance as yf\n",
    "\n",
    "class StockState(TypedDict):\n",
    "    ticker: str\n",
    "\n",
    "def fetch_stock_details(state: StockState):\n",
    "    period = \"1mo\"\n",
    "    stock_symbol = state[\"ticker\"]\n",
    "    try:\n",
    "        stock = yf.Ticker(stock_symbol)\n",
    "\n",
    "        # Retrieve general stock info and historical market data\n",
    "        stock_info = stock.info  # Basic company and stock data\n",
    "        stock_history = stock.history(period=period).to_dict()  # Historical OHLCV data\n",
    "\n",
    "        # Combine both into a single dictionary\n",
    "        stock = {\n",
    "            \"stock_symbol\": stock_symbol,\n",
    "            \"info\": stock_info,\n",
    "            \"history\": stock_history\n",
    "        }\n",
    "        \n",
    "        return {\"stock_details\": [str(stock)]}\n",
    "        \n",
    "    except Exception as e:\n",
    "        return {\"error\": f\"Error fetching stock data for {stock_symbol}: {str(e)}\"}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = fetch_stock_details({\n",
    "    \"ticker\": \"AI\"\n",
    "})\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analyze all the results (Reduce Step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "def generate_stock_recommendations(state: InvestmentAdvisorState):\n",
    "    financial_data = \"\\n\\n\\n\".join(state[\"stock_details\"])\n",
    "\n",
    "    prompt = f\"\"\"\n",
    "    You are a professional investment advisor.\n",
    "\n",
    "    Below is a list of companies with their financial data:\n",
    "\n",
    "    {financial_data}\n",
    "\n",
    "    Your task is to:\n",
    "    1. Analyze each company based on the provided technical information (consider only companies from the list above).\n",
    "    2. Write a short paragraph for each company that includes:\n",
    "    - The company name and ticker\n",
    "    - A brief description\n",
    "    - Why it may or may not be a good investment\n",
    "    3. Rank the companies from best to worst in terms of investment potential.\n",
    "    4. Return the result as a sorted list (highest priority first), where each item includes:\n",
    "    - Rank (starting from 1)\n",
    "    - Ticker\n",
    "    - Company name\n",
    "    - Short description\n",
    "    - Reason for its ranking\n",
    "    \"\"\"\n",
    "\n",
    "    recommendation = llm.invoke([HumanMessage(content=prompt)])\n",
    "\n",
    "    return {\"recommendation\": recommendation.content}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mapping (Fetch Stock Details)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "from langgraph.graph import END, StateGraph, START\n",
    "\n",
    "builder = StateGraph(InvestmentAdvisorState)\n",
    "\n",
    "builder.add_node(\"generate_list_of_stocks\", generate_list_of_stocks)\n",
    "builder.add_node(\"fetch_stock_details\", fetch_stock_details)\n",
    "builder.add_node(\"generate_stock_recommendations\", generate_stock_recommendations)\n",
    "\n",
    "\n",
    "builder.add_edge(START, \"generate_list_of_stocks\")\n",
    "\n",
    "\n",
    "# Map - Reducing\n",
    "from langgraph.constants import Send\n",
    "def continue_to_details(state: InvestmentAdvisorState):\n",
    "    return [Send(\"fetch_stock_details\", {\"ticker\": ticker}) for ticker in state[\"stock_tickers\"]]\n",
    "\n",
    "builder.add_conditional_edges(\"generate_list_of_stocks\", continue_to_details, [\"fetch_stock_details\"])\n",
    "\n",
    "\n",
    "\n",
    "builder.add_edge(\"fetch_stock_details\", \"generate_stock_recommendations\")\n",
    "builder.add_edge(\"generate_stock_recommendations\", END)\n",
    "\n",
    "graph = builder.compile()\n",
    "Image(graph.get_graph().draw_mermaid_png())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = graph.invoke({\"financial_area\": \"small AI companies\", \"stock_number\": 5})\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(result[\"recommendation\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
